#    Copyright (C) 2003-2008 William Waites
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""
Implementation of elements of an Autonomous System Graph G(V,E) where
Vertexes are AS objects and Edges can be derived from the paths contained
in BGP updates.
"""
from radix import Radix
import networkx as NX
from dpkt import bgp
import socket
import bsddb

__all__ = ["AS", "Path", "Mesh"]

class AS(object):
	"""
	Autonomous System Object

	This object has some special properties, namely that the
	constructor will return *the same* object if the argument
	(the ASN) is the same. The object is fully pickleable as
	well.

	>>> a = AS(123)
	>>> print a
	AS123
	>>>
	"""
	__networks__ = {}
	@staticmethod
	def __new__(cls, asn, *av, **kw):
		try:
			return cls.__networks__[asn]
		except KeyError:
			obj = super(AS, cls).__new__(cls, asn)
			cls.__networks__[asn] = obj
			return obj
	def __init__(self, asn):
		self.asn = asn
		self.__dict__.setdefault("prefixes", [])
	def __reduce__(self):
		"""
		To verify that the pickle interface as well as the 
		constructor work correctly, construct an object and
		pickle then unpickle it. We should obtain the very
		same object.

		>>> import pickle
		>>> a = AS(123)
		>>> b = pickle.loads(pickle.dumps(a))
		>>> a is b
		True
		>>>
		"""
		return AS, (self.asn,), {"prefixes": self.prefixes}
	def __str__(self):
		return "AS%s" % (self.asn,)
	def __repr__(self):
		return "AS(%s)" % (self.asn,)
	def announce(self, route):
		"""
		Add a route to the list of prefixes announced by this
		autonomous system.

		>>> a = AS(123)
		>>> a.announce(("192.0.2.0", 24))
		>>> a.announce(("192.0.2.1", 32))
		>>> print a.prefixes
		[('192.0.2.0', 24), ('192.0.2.1', 32)]
		>>>
		"""
		if route not in self.prefixes:
			self.prefixes.append(route)
	def withdraw(self, route):
		if route in self.prefixes:
			self.prefixes.remove(route)
	def size(self, af=socket.AF_INET):
		"""
		Size takes an optional address family argument and returns
		the aggregated size of the autonomous system (number of 
		addresses). Note that this example will use the prefixes
		added in the doctest for the announce method.

		>>> a = AS(123)
		>>> print a.size(socket.AF_INET)
		256
		>>>
		"""
		size = 0
		rtree = Radix()
		if af == socket.AF_INET:
			afbits = 32
		elif af == socket.AF_INET6:
			afbits = 128

		for net, mask in self.prefixes:
			rtree.add(network=net, masklen=mask)
		def _routes():
			for r in rtree:
				if r.family == af:
					yield r
		routes = list(_routes())
		routes.sort(lambda x,y: cmp(y.prefixlen,x.prefixlen))

		for node in routes:
			if node.family != af:
				continue
			rtree.delete(node.prefix)
			if rtree.search_best(node.network) is None:
				size = size + 2**(afbits-node.prefixlen)
		return size

class Path(list):
	"""
	The Path class represents an Autonomous System Path, it uses the 
	ASPath class from dpkt.bgp but provides a simple list interface,
	with the elements being AS objects or lists of AS objects in the
	case of an AS_SET. This class is pickleable, saving the binary
	representation of the path and reconstituting itself therefrom.

	Like AS the constructor will return an existing instance of a path
	if it has been seen before rather than constructing a distinct 
	object with the (by definition) same data. This allows for efficient
	in-core storage.

	>>> import pickle
	>>> data = '\\x02\\x06\\x0c\\xcd\\x0c\\xc3\\x00\\xae\\x00\\xef\\x00\\xef\\x00\\xef'
	>>> p = Path(data)
	>>> print p
	[AS(3277), AS(3267), AS(174), AS(239), AS(239), AS(239)]
	>>> q = pickle.loads(pickle.dumps(p))
	>>> p is q
	True
	>>>
	"""
	__paths__ = {}
	@staticmethod
	def __new__(cls, path):
		path_s = str(path)
		try:
			return cls.__paths__[path_s]
		except KeyError:
			obj = super(Path, cls).__new__(cls, path)
			obj.__init__(path)
			cls.__paths__[path_s] = obj
			return obj
	def __init__(self, path):
		if not self:
			if isinstance(path, str):
				path = bgp.BGP.Update.Attribute.ASPath(path)
			self.data = path
			for seg in path.segments:
				if seg.type == bgp.AS_SEQUENCE:
					self.extend(map(AS, seg.path))
				elif seg.type == bgp.AS_SET:
					self.append(map(AS, seg.path))
	def __reduce__(self):
		return Path, (str(self.data),)
	def __getstate__(self):
		return False
		
	def __hash__(self):
		return hash(str(self.data))

def _xpermutations(items, n):
	"""
	Utility function to return (order independent) permutations of
	a sequence. This is used to construct a fully-interconnected
	mesh from a list, when used with n=2. This function is not exported
	from the module.

	>>> list(_xpermutations([1,2,3,4], 2))
	[[1, 2], [1, 3], [1, 4], [2, 3], [2, 4], [3, 4]]
	>>>
	"""
	if n==0: yield []
	else:
		for i in xrange(len(items)):
			for cc in _xpermutations(items[i+1:],n-1):
				yield [items[i]]+cc

class Mesh(NX.Graph):
	"""
	A derivitive of NX.Graph with support for some specialized aspects of
	the graph of autonomous systems on the Internet.
	"""
	def __init__(self, *av, **kw):
		super(Mesh, self).__init__(*av, **kw)
		self.__paths__ = bsddb.hashopen(None, "c")
	def __reduce__(self):
		f,a,d = super(Mesh, self).__reduce__()
		e = d.copy()
		del e["__paths__"]
		return f,a,e
	def __setstate__(self, state):
		self.__dict__.update(state)
		self.__paths__ = bsddb.hashopen(None, "c")
	def add_path(self, path):
		"""
		This add_path method optimizes for the case where (1) the
		path argument is an instance of graph.Path and (2) it has already
		been added to the graph (in which case nothing is done). This is
		because when processing BGP updates we see the same path many times
		and re-adding the path elements (which is essentially a noop) becomes
		needlessly computationally intensive.

		Of course a simple list or iterable can be passed in as with the
		parent class' implementation.

		>>> m = Mesh()	
		>>> p1 = '\\x02\\x06\\x0c\\xcd\\x0c\\xc3\\x00\\xae\\x00\\xef\\x00\\xef\\x00\\xef'
		>>> p2 = '\\x02\\x05\\x01\\x1e\\x00\\xae\\x00\\xef\\x00\\xef\\x00\\xef'
		>>> p3 = '\\x02\\x053\\xb5\\x00\\xae\\x00\\xef\\x00\\xef\\x00\\xef'
		>>> m.add_path(Path(p1))
		>>> m.add_path(Path(p2))
		>>> m.add_path(Path(p3))
		>>> n = m.neighbors(AS(174))
		>>> n.sort()
		>>> print n
		[AS(286), AS(13237), AS(3267), AS(239)]
		>>>
		"""
		if isinstance(path, Path):
			data = str(path.data)
			if data in self.__paths__:
				return
			else:
				self.__paths__[data] = None
		super(Mesh, self).add_path(path)	
	def add_edge(self, u, v):
		"""
		If either node u or node v is a list (such as is the case with an
		AS_SET) we treat them as internally fully meshed and create the
		appropriate links. The other node is then connected to all of them.

		For example,

		>>> m = Mesh()
		>>> m.add_edge(AS(1), [AS(2), AS(3), AS(4)])
		>>> print m.neighbors(AS(1))
		[AS(2), AS(3), AS(4)]
		>>> print m.neighbors(AS(2))
		[AS(1), AS(3), AS(4)]
		>>>
		"""
		if isinstance(u, list):
			for w in u: self.add_edge(w,v)
			for w,x in _xpermutations(u,2): self.add_edge(w,x)
		elif isinstance(v, list):
			for w in v: self.add_edge(u,w)
			for w,x in _xpermutations(v,2): self.add_edge(w,x)
		else:
			super(Mesh, self).add_edge(u,v)

if __name__ == '__main__':
	import doctest
	doctest.testmod()
